import util.lib as lib

import pandas as pd
from datetime import datetime
from dateutil.relativedelta import relativedelta

from month import cal_month_inc_num
from week import cal_week_inc_num

import util.mydecorator

from daily import daily_share_file_exists as daily_share_file_exists
from daily import get_daily_df

# 数据按周分组
def lib_format_week_date(x):
    week_start = x - relativedelta(days=x.weekday())
    week_end = week_start + relativedelta(days=6)
    return x.strftime('%Y-第%W周') +  ':' +  week_start.strftime('%m%d') + '--' +  week_end.strftime('%m%d') 

def format_week(df, stock_name):
    # 每周末 level_point
    tmp = df[['rel_lp']].groupby(lib_format_week_date).agg(lambda x: x.iloc[-1]).copy()
    
    tmp_level_point = df[['level_point']].groupby(lib_format_week_date).agg(lambda x: x.iloc[-1])
    tmp['level_point'] = tmp_level_point['level_point']

    tmp_drawdown = df[['drawdown']].groupby(lib_format_week_date).agg(lambda x: x.iloc[-1])
    tmp['drawdown'] = tmp_drawdown['drawdown']
    
    tmp_close = df[['close']].groupby(lib_format_week_date).agg(lambda x: x.iloc[-1])
    tmp['close'] = tmp_close['close']
    
    week_pct = 100 * (tmp_close['close'] - tmp_close['close'].shift()) / tmp_close['close'].shift()
    tmp['week_pct'] = round(week_pct,2)
    
    df_week_gb = tmp['close']
    
    tmp['15w_gb'] = df_week_gb.rolling(window=15).mean()
    tmp['15w_gb'] = round(tmp['15w_gb'], 3)

    tmp['15dis_pct'] = round(( df_week_gb/tmp['15w_gb'] -1) * 100, 2)

    tmp['20w_gb'] = df_week_gb.rolling(window=20).mean()
    tmp['20w_gb'] = round(tmp['20w_gb'], 3)
    tmp['20dis_pct'] = round(( df_week_gb/tmp['20w_gb'] -1) * 100, 2)

    tmp['30w_gb'] = df_week_gb.rolling(window=30).mean()
    tmp['30w_gb'] = round(tmp['30w_gb'], 3)
    tmp['30dis_pct'] = round(( df_week_gb/tmp['30w_gb'] -1) * 100, 2)

    tmp['60w_gb'] = df_week_gb.rolling(window=60).mean()
    tmp['60w_gb'] = round(tmp['60w_gb'], 3)    

    tmp['dis_pct'] = round(( df_week_gb/tmp['60w_gb'] -1) * 100, 2)
    
    tmp = tmp.dropna()
    if len(tmp) <= 0:
        return pd.DataFrame()

    # pb & pe
    # 防止最新的数据没更新
    #df[['pb_ratio']] = df['pb_ratio'].fillna(method='ffill')
    tmp_pb = df[['pb_ratio']].groupby(lib_format_week_date).agg(lambda x: x.iloc[-1])
    tmp['pb_ratio'] = tmp_pb['pb_ratio']

    #df[['pe_ratio']] = df['pe_ratio'].fillna(method='ffill')
    tmp_pe = df[['pe_ratio']].groupby(lib_format_week_date).agg(lambda x: x.iloc[-1])
    tmp['pe_ratio'] = tmp_pe['pe_ratio']

    tmp["stock"] = stock_name

    df = tmp

    df['inc_num'] = 0
    df['pct_sum'] = 0.0
    df = cal_week_inc_num(df)

    return df[['stock','close','week_pct','pct_sum','inc_num','drawdown','level_point','rel_lp',
                '15w_gb','20w_gb','30w_gb', '60w_gb',
                '15dis_pct', '20dis_pct', '30dis_pct','dis_pct',
                'pb_ratio','pe_ratio'
               ]][-450:]



def format_month(df, stock_name):
    
    # 月末 level_point
    tmp = df[['rel_lp']].groupby(lambda x: x.strftime('%Y--%m')).agg(lambda x: x.iloc[-1]).copy()
    
    tmp_level_point = df[['level_point']].groupby(lambda x: x.strftime('%Y--%m')).agg(lambda x: x.iloc[-1])
    tmp['level_point'] = tmp_level_point['level_point']

    tmp_drawdown = df[['drawdown']].groupby(lambda x: x.strftime('%Y--%m')).agg(lambda x: x.iloc[-1])
    tmp['drawdown'] = tmp_drawdown['drawdown']

    tmp_gt1200d = df[['gt1200d']].groupby(lambda x: x.strftime('%Y--%m')).agg(lambda x: x.iloc[-1])
    tmp['gt1200d'] = round(tmp_gt1200d['gt1200d'], 3)
    
    tmp_close = df[['close']].groupby(lambda x: x.strftime('%Y--%m')).agg(lambda x: x.iloc[-1])
    tmp['close'] = tmp_close['close']
    
    month_pct = 100 * (tmp_close['close'] - tmp_close['close'].shift()) / tmp_close['close'].shift()
    tmp['month_pct'] = round(month_pct,2)

    # 最低价格
    tmp_low_close = df[['close']].groupby(lambda x: x.strftime('%Y--%m')).agg(lambda x: x.min()) 
    tmp['low_close'] = tmp_low_close['close']

    month_low_pct = 100 * (tmp['low_close'] - tmp_close['close'].shift()) / tmp_close['close'].shift()
    tmp['low_month_pct'] = round(month_low_pct, 2)
    
    tmp = tmp.dropna(subset=['month_pct'])
    
    tmp["stock"] = stock_name

    # 涨跌统计
    df = tmp

    df['inc_num'] = 0
    df['pct_sum'] = 0.0
    df = cal_month_inc_num(df)

    return df[['stock','close','month_pct','pct_sum','inc_num',"drawdown",'level_point','rel_lp','gt1200d','low_close','low_month_pct']][-140:]
    
def task(stock_name, month_writer, week_writer):

    if daily_share_file_exists(stock_name):
        df = get_daily_df(stock_name)

        # 2. 拼装数据
        month_df = format_month(df, stock_name)

        week_df = format_week(df, stock_name)
            
        # 3. 写入 excel
        if len(month_df) > 0:
            month_df.to_excel(month_writer, sheet_name=stock_name)
        if len(week_df) > 0:
            week_df.to_excel(week_writer,sheet_name=stock_name)

        print(stock_name, " finished")


def save_to_excel(shares, date):
    n = 0  
    print("统计时间：", date)

    month_writer = pd.ExcelWriter(lib.data_path + 'month_level_point.xlsx')
    week_writer = pd.ExcelWriter(lib.data_path + 'week_level_point.xlsx')

    for item in shares.iterrows():
        stock = item[1]['trade_code']
        stock_name = item[1]['name']
        print(stock, stock_name)
        task(stock_name, month_writer, week_writer)

    month_writer.close()
    week_writer.close()
    print("结束")

@util.mydecorator.calTime
def job():

    date = datetime.now()
    df = lib.lib_get_all_stock()
    save_to_excel(df, date)

if __name__ == '__main__':
    job()